import argparse
import glob
import h5py
import numpy as np
import PIL.Image as pil_image
from utils import convert_rgb_to_y
import os


def train(args):
    h5_file = h5py.File(args.output_path, 'w')

    hr_patches = []
    lr_patches = []

    for hr_image_path in sorted(glob.glob('{}/*'.format(args.hr_images_dir))):
        # 将HR和LR对应文件夹下的文件同时读入
        (filepath, filename) = os.path.split(hr_image_path)
        (root_path, hr_path) = os.path.split(filepath)
        lr_path = 'LR'
        lr_image_path = os.path.join(root_path, lr_path, filename)
        # ***************************** #

        hr = pil_image.open(hr_image_path).convert('RGB')
        hr = hr.resize((hr.width, hr.height))
        hr = np.array(hr).astype(np.float32)  # BGR颜色，所以是3通道图片
        hr = convert_rgb_to_y(hr)

        lr = pil_image.open(lr_image_path).convert('RGB')
        lr = lr.resize((lr.width, lr.height))
        lr = np.array(lr).astype(np.float32)
        lr = convert_rgb_to_y(lr)


        for i in range(0, lr.shape[0] - args.patch_size + 1, args.stride):
            for j in range(0, lr.shape[1] - args.patch_size + 1, args.stride):
                lr_patches.append(lr[i:i + args.patch_size, j:j + args.patch_size])
                hr_patches.append(hr[i:i + args.patch_size, j:j + args.patch_size])

    lr_patches = np.array(lr_patches)
    hr_patches = np.array(hr_patches)

    h5_file.create_dataset('lr', data=lr_patches)
    h5_file.create_dataset('hr', data=hr_patches)

    h5_file.close()


def eval(args):
    h5_file = h5py.File(args.output_path, 'w')

    lr_group = h5_file.create_group('lr')
    hr_group = h5_file.create_group('hr')

    for i, image_path in enumerate(sorted(glob.glob('{}/*'.format(args.hr_images_dir)))):
        # 将HR和LR对应文件夹下的文件同时读入
        (filepath, filename) = os.path.split(image_path)
        (root_path, hr_path) = os.path.split(filepath)
        lr_path = 'LR'
        lr_image_path = os.path.join(root_path, lr_path, filename)
        # ***************************** #
        hr = pil_image.open(image_path).convert('RGB')
        hr = hr.resize((hr.width, hr.height))
        hr = np.array(hr).astype(np.float32)
        hr = convert_rgb_to_y(hr)
        hr_group.create_dataset(str(i), data=hr)

        lr = pil_image.open(lr_image_path).convert('RGB')
        lr = lr.resize((lr.width, lr.height))
        lr = np.array(lr).astype(np.float32)
        lr = convert_rgb_to_y(lr)
        lr_group.create_dataset(str(i), data=lr)


    h5_file.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # parser.add_argument('--images-dir', type=str, default=r'D:\codes\01-rebuild\SR\FSRCNN-master\dataset\train\91',required=True)   ---lcj
    # parser.add_argument('--output-path', type=str, default=r'D:\codes\01-rebuild\SR\20220211j7in0LBK\SRCNN-pytorch-master\train-lcj.h5 ',  required=True)
    # parser.add_argument('--hr_images-dir', type=str, default=r'C:\lcj\02-dataset\6.18_zerPSF_new\dataset_hr_lr2\HR')
    parser.add_argument('--hr_images-dir', type=str, default=r'hr文件路径')

    parser.add_argument('--output-path', type=str,
                        default=r'权重保存位置')
    parser.add_argument('--patch-size', type=int, default=32)  # lcj -- 之前是33
    parser.add_argument('--stride', type=int, default=14)
    args = parser.parse_args()

    train(args)
    # eval(args)
    # if  not args.eval:   # 没使用“评估”就是 train（）
    #     train(args)
    # else:                # 使用“评估”就是eval() ___lcj
    #     eval(args)